#!/usr/bin/env python
# coding: utf-8
import os
import optparse
import select
import socket
import struct
import sys
from cStringIO import StringIO

DEFAULT_HOST = '127.0.0.1'
DEFAULT_PORT = 0x2304

CHUNK_COMMAND = 0
CHUNK_STDIN = 1
CHUNK_STDOUT = 2
CHUNK_STDERR = 3
CHUNK_EXIT = 4

def write_str(out, data):
    out.write(struct.pack('!H%ds' % len(data), len(data), data))

def read_str(sock):
    count = struct.unpack('!H', sock.recv(2))[0]
    return sock.recv(count)

def write_cmd(out, args):
    cmdname = args[0]
    cmdargs = args[1:]
    cwd = os.path.abspath('.')
    out.write(chr(CHUNK_COMMAND))
    write_str(out, cmdname)
    out.write(struct.pack('!H', len(cmdargs)))
    for arg in cmdargs:
        write_str(out, arg)
    write_str(out, cwd)

class Struct(object):
    def __init__(self, **kwargs):
        for name, value in kwargs.items():
            setattr(self, name, value)

ST_CHUNK_BEGIN = 0
ST_CHUNK_TYPE = 1
ST_CHUNK_HEADER = 2
class ChunkDecoder(object):
    def __init__(self):
        self.buf = ''
        self.state = ST_CHUNK_BEGIN
        self.chunk = Struct()       # current chunk

    def decode(self, data):
        self.buf += data
        chunks = []
        pos = 0
        while True:
            if self.state == ST_CHUNK_BEGIN:
                if len(self.buf) >= pos+1:
                    self.state = ST_CHUNK_TYPE
                    self.chunk.type = ord(self.buf[pos])
                    pos += 1
                else:
                    self.buf = self.buf[pos:]
                    break
            elif self.state == ST_CHUNK_TYPE:
                if len(self.buf) >= pos+2:
                    self.state = ST_CHUNK_HEADER
                    self.chunk.length = struct.unpack('!H', self.buf[pos:pos+2])[0]
                    pos += 2
                else:
                    self.buf = self.buf[pos:]
                    break
            else:   # ST_CHUNK_HEADER
                if len(self.buf) >= pos+self.chunk.length:
                    self.state = ST_CHUNK_BEGIN
                    self.chunk.value = self.buf[pos:pos+self.chunk.length]
                    pos += self.chunk.length
                    chunks.append(self.chunk)
                    self.chunk = Struct()
                else:
                    self.buf = self.buf[pos:]
                    break
        return chunks

if __name__ == '__main__':
    parser = optparse.OptionParser(add_help_option=False, usage='%prog [options] -- command [args]')
    parser.add_option('--help', action='store_true', help='show this help message and exit', default=False)
    parser.add_option('-h', '--host', dest='host', help='specify the server host to connect', default=DEFAULT_HOST)
    parser.add_option('-p', '--port', dest='port', type='int', help='specify the server port to connect, default 0x2304', default=DEFAULT_PORT)
    opts, args = parser.parse_args()

    if opts.help or len(args) <= 0:
        parser.print_help()
        sys.exit(1)

    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.connect((opts.host, opts.port))
    out = StringIO()
    write_cmd(out, args)
    sock.send(out.getvalue())

    chunk_decoder = ChunkDecoder()

    eof = False
    while not eof:
        rlist, wlist, xlist = select.select([sys.stdin, sock], [], [])
        for r in rlist:
            if r == sys.stdin:
                data = sys.stdin.readline()
                buf = StringIO()
                buf.write(chr(CHUNK_STDIN))
                buf.write(struct.pack('!H', len(data)))
                buf.write(data)
                sock.send(buf.getvalue())
            elif r == sock:
                data = sock.recv(1024)
                if not data:    # connection closed
                    eof = True
                    break
                chunks = chunk_decoder.decode(data)
                for chunk in chunks:
                    if chunk.type == CHUNK_STDOUT:
                        sys.stdout.write(chunk.value)
                        sys.stdout.flush()
                    elif chunk.type == CHUNK_STDERR:
                        sys.stderr.write(chunk.value)
                    elif chunk.type == CHUNK_EXIT:
                        sys.exit(ord(chunk.value))

